#
# messkit.stream
#
# Copyright 2004-2006 Helsinki Institute for Information Technology (HIIT)
# and the authors.
#
# Authors: Ken Rimey <rimey@hiit.fi>
#

# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

"""
Bidirectional byte stream with read and write buffering

This module defines three classes: Stream, StreamServer, and
StreamManager.  Instances of Stream are connection endpoints,
while instances of StreamServer listen for incoming connections.

These are abstract base classes.  See the available implementation
modules (*_stream.py) for instantiable classes.
"""

from socket import AF_INET, SOCK_STREAM
from traceback import print_exc

from messkit.buffer import Buffer

class Stream:
    """
    Connection endpoint
    """
    flush_size = 65536
    send_size = 8192
    recv_size = 16384          # Minimum of 16384 required by tlslite.

    def __init__(self, mgr, socket=(AF_INET, SOCK_STREAM), address=None):
        """
        Initialize a stream instance.

        To initialize the initiating endpoint of a stream, provide a
        reference to a StreamManager instance and optionally a tuple
        of parameters for initializing a socket.

        Initialization of the non-initiating endpoint of a stream is
        done in an accept callback, using the callback arguments to
        initialize the Stream instance.  In this case, the "socket"
        parameter will actually be a socket rather than a tuple of
        parameters.
        """
        mgr.streams.append(self)
        self.mgr = mgr
        self.initiating = isinstance(socket, tuple)
        self.connected = not self.initiating
        assert bool(address) == self.connected
        self.address = address
        self.eof_seen = False
        self.closed = False
        self.error = None
        self.error_handler = None
        self.rbuf = Buffer()
        self.wbuf = Buffer()
        self.bytes_in = 0
        self.bytes_out = 0
        self._init(socket)

    def set_error_handler(self, cb):
        """
        Set up an error handler

        The error handler is invoked with an argument providing
        further information (a tuple or an exception object) about an
        error signaled by a send or recv socket operation.  All such
        errors are reported through this mechanism, and never via
        stream methods raising exceptions.  Errors in connecting the
        socket are reported through the connect callback, however.

        Once an error has occurred, fill callbacks will generally no
        longer be invoked, and any socket operations attempted will
        simply have no effect.  You still need to call the close()
        method, and its callback will be invoked as usual.
        """
        self.error_handler = cb
        if self.error:
            self.issue_callback(self.error_handler, self.error)

    def close(self, cb=None):
        """
        Close the stream.

        This triggers a flush and the eventual freeing of all
        underlying resources.  Always call this method on an instance
        that you are finished with, even if it is not yet connected,
        or an error has occurred.  It is never invoked automatically.
        It can be called more than once.

        The optional callback is invoked with no arguments when the
        underlying socket is actually closed.  It is only invoked on
        the first call to close().

        If an error occurs in the implicit flush, the error handler
        and the close callback will both be invoked, in that order.
        Note that if the error handler calls close(), the close
        callback that is invoked will be the one specified in the
        original call to close(), not the one specified in the second
        call to close() in the error handler.
        """
        def proxy():
            self.mgr.streams.remove(self)
            self.error_handler = None # Avoid cycles.
            if cb is not None:
                cb()

        if not self.closed:
            self.closed = True
            if self.wbuf.size:
                self._flush(self.wbuf)
            self._close(proxy)

    def connect(self, address, cb=None):
        """
        Connect the stream to a remote address.

        The argument is a tuple (host, port) specifying the remote
        address to connect to.

        This always returns immediately.  The callback will be invoked
        with an argument detailing any error that may have occurred.
        This value will test false if the connection was established
        successfully.

        Do not call this method more than once, and do not call it on
        an instance returned by accept(), which will already be connected.
        """
        def proxy(error):
            self.connected = not error
            if cb is not None:
                cb(error)

        self._check_if_closed()
        self.address = address
        self._connect(address, proxy)

    def write(self, data):
        """
        Append a string to the output buffer.

        A flush is triggered when the buffer reaches or exceeds a
        certain size.
        """
        self._check_if_closed()
        self.wbuf.append(data)
        if self.wbuf.size >= self.flush_size:
            self._flush(self.wbuf)

    def flush(self):
        """
        Flush the output buffer.

        This triggers the asynchronous transmission of the entire
        contents of the output buffer.
        """
        self._check_if_closed()
        if self.wbuf.size:
            self._flush(self.wbuf)

    def __len__(self):
        """
        Return the number of bytes present in the input buffer.
        """
        return self.rbuf.size

    def peek(self, size, start=0):
        """
        Return at most "size" bytes, leaving them in the input buffer.

        If "start" is specified, it is an offset starting from which
        the requested number of bytes is returned.
        """
        return self.rbuf.peek(size, start)

    def skip(self, size):
        """
        Read and discard at most "size" bytes from the input buffer.
        """
        self.rbuf.skip(size)

    def read(self, size):
        """
        Read at most "size" bytes from the input buffer.

        This only returns as many bytes as are available in the buffer.
        """
        return self.rbuf.read(size)

    def fill(self, cb=None, size=0):
        """
        Trigger asynchronous filling of the input buffer.

        The "size" parameter is the target for the number of bytes in
        the input buffer.  If it is not specified, at least one byte
        (typically more) is added to the buffer, unless EOF is seen
        immediately.  This is also the behavior if there are already
        "size" bytes in the buffer.  Otherwise, the number of
        additional bytes required is calculated, and filling continues
        until at least that many bytes are added or EOF is seen.

        The callback is invoked when the job is done, passing the
        number of bytes added to the buffer.  It is an error to call
        this method again until the callback is received.

        If the argument to the callback is zero (or if it is less than
        the number of bytes requested), one can conclude that EOF has
        been read and further calls to fill() will have no effect.
        The stream remains open, however, and it may be possible to
        continue transmitting data.
        """
        self._check_if_closed()
        self._fill(self.rbuf, max(1, size - self.rbuf.size), cb)

    def get_socket(self):
        self._check_if_closed()
        return self._get_socket()

    socket = property(get_socket)

    def _init(self, socket):
        raise NotImplementedError

    def _close(self, cb):
        raise NotImplementedError

    def _connect(self, address, cb):
        raise NotImplementedError

    def _flush(self, buf):
        raise NotImplementedError

    def _fill(self, buf, size, cb):
        raise NotImplementedError

    def _get_socket(self):
        raise NotImplementedError

    def _check_if_closed(self):
        if self.closed:
            raise ValueError, 'Operation on closed stream.'

    def _signal_error(self, error):
        # XXX Should really be made thread-safe.
        assert error
        if not self.error:
            self.error = error
            self.issue_callback(self.error_handler, error)

    def issue_callback(self, f, *args):
        self.mgr.issue_main_thread_callback(f, *args)

class StreamServer:
    """
    Passive endpoint to listen for new connections
    """
    def __init__(self, mgr, socket_params=(AF_INET, SOCK_STREAM)):
        """
        Initialize a stream server instance.

        The first argument is a reference to a StreamManager instance.
        The optional second argument is a tuple of parameters for
        initializing a listening socket.
        """
        mgr.servers.append(self)
        self.mgr = mgr
        self.closed = False
        self._init(socket_params)

    def close(self, cb=None):
        """
        Shut down the stream server.

        Always call this method on an instance that you are finished
        with.  It can be called more than once.

        The optional callback is invoked with no arguments when the
        underlying socket is actually closed.  It is only called on
        the first call to close().

        This does not close previously created streams, on the
        presumption that it is more natural for the caller to keep
        track of them.
        """
        def proxy():
            self.mgr.servers.remove(self)
            if cb is not None:
                cb()

        if not self.closed:
            self.closed = True
            self._close(proxy)

    def listen(self, address):
        """
        Bind the stream server to a local address and begin listening.

        The argument is a tuple specifying the local address to bind
        to, or other parameters for initializing the server.
        """
        if self.closed:
            raise ValueError, 'Stream server has been closed.'

        self._listen(address)

    def accept(self, cb):
        """
        Accept one incoming connection.

        This always returns immediately.  The callback will be invoked
        with arguments suitable for initializing a new Stream instance.
        """
        def proxy(conn, addr):
            cb(self.mgr, conn, addr) # Positional arguments only.

        if self.closed:
            raise ValueError, 'Operation on closed stream server.'

        self._accept(proxy)

    def _init(self, socket_params):
        raise NotImplementedError

    def _close(self, cb):
        raise NotImplementedError

    def _listen(self, address):
        raise NotImplementedError

    def _accept(self, cb):
        raise NotImplementedError

    def issue_callback(self, f, *args):
        self.mgr.issue_main_thread_callback(f, *args)

class StreamManager:
    """
    Manager for streams and servers sharing a common implementation
    """
    def __init__(self, callback_callback=None):
        """
        Initialize the stream manager.

        The callback takes one argument, which will be a function
        taking no arguments.  The client must arrange for this
        function to be invoked in the program's main thread, no matter
        what thread invoked the callback.  This protocol makes it
        possible for other callbacks to be invoked by the main thread.
        """
        self.callback = callback_callback
        self.closed = False
        self.streams = []
        self.servers = []

    def create_stream(self, *args):
        """
        Stream factory

        This provides an alternative to directly initializing an
        instance of the appropriate Stream implementation class.
        """
        return self._create_stream(*args)

    def create_server(self, socket_params=(AF_INET, SOCK_STREAM)):
        """
        StreamServer factory

        This provides an alternative to directly initializing an
        instance of the appropriate StreamServer implementation class.
        """
        return self._create_server(socket_params)

    def close(self, cb=None):
        """
        Release any resources associated with the manager.

        The optional callback is invoked with no arguments in the
        main thread when the manager has finished shutting down.
        """
        if not self.closed:
            self.closed = True

            for stream in list(self.streams):
                stream.close()

            for server in list(self.servers):
                server.close()

            # An alternative would be to specify a callback when
            # closing the open streams and servers above, and to
            # execute the following in the callback when self.streams
            # and self.servers are both empty.
            self._close(cb)

    def _create_stream(self, socket):
        raise NotImplementedError

    def _create_server(self, socket_params):
        raise NotImplementedError

    def _close(self, cb):
        self.issue_main_thread_callback(cb)

    def issue_main_thread_callback(self, f, *args):
        def proxy():
            try:
                f(*args)
            except:
                print_exc()

        cb = self.callback
        if f is not None and cb is not None:
            cb(proxy)
